from collections import defaultdict
import click
import pickle as pk
import os
from loguru import logger


@click.command()
@click.option("-p", "--path", help="The dir path of pkl files")
def run(path):
    if not os.path.exists(path):
        logger.error("Path Not Found!")
        exit(0)

    for file in os.listdir(path):
        if "pkl" not in file:
            continue
        with open("{}/{}".format(path, file), "rb") as f:
            data = pk.load(f)
            logger.debug("Processing {} ...".format(file))
            for funcname in data:
                bblock = data[funcname]["blocks"]
                edges = data[funcname]["edges"]
                inline_addr_bb_range = data[funcname]["inline_addr_bb_range"]
                inline_addr_from_idx = data[funcname]["inline_addr_from_idx"]
                indegree_arr = [0 for _ in range(len(bblock))]
                outdegree_arr = [0 for _ in range(len(bblock))]
                non_inline_end_pos = len(bblock) - 1 if len(inline_addr_bb_range) == 0 else inline_addr_bb_range[0][0]

                delete_edges = []
                add_edges = []
                # Clean the retn arrows
                for i in range(len(edges)):
                    fromb, tob = edges[i]
                    if fromb <= non_inline_end_pos and "retn" in bblock[fromb]:
                        delete_edges.append([fromb, tob])
                        continue
                    indegree_arr[tob] += 1
                    outdegree_arr[fromb] += 1

                inline_end_dict = defaultdict(lambda: [])

                # If we want to extend bblock "B"
                # We need to let "B" -> "C" while "C" is the inline bblock from "B" and its indegree is 0.
                for i in range(len(bblock)):
                    if indegree_arr[i] == 0 and i > non_inline_end_pos:
                        for ori_bb_pos, bb_range in zip(inline_addr_from_idx, inline_addr_bb_range):
                            if i > bb_range[0] and i <= bb_range[1]:
                                add_edges.append([ori_bb_pos, i])
                                break
                    if outdegree_arr[i] == 0 and i > non_inline_end_pos:
                        for ori_bb_pos, bb_range in zip(inline_addr_from_idx, inline_addr_bb_range):
                            if i > bb_range[0] and i <= bb_range[1]:
                                inline_end_dict[ori_bb_pos].append(i)
                                break

                # Change the whole graph
                for i in range(len(edges)):
                    fromb, tob = edges[i]
                    if fromb in inline_addr_from_idx:
                        delete_edges.append([fromb, tob])
                        for endpos in inline_end_dict[fromb]:
                            add_edges.append([endpos, tob])
                logger.debug(funcname)
                logger.debug(edges)
                logger.debug(add_edges)
                logger.debug(delete_edges)
                for x in add_edges:
                    edges.append(x)
                for x in delete_edges:
                    edges.remove(x)

                data[funcname]["edges"] = edges

            with open("{}/{}.new".format(path, file), "wb") as f:
                pk.dump(data, f)


if __name__ == "__main__":
    run()
